"""音频特征提取"""

import torch
import torch.nn as nn
from tqdm import tqdm
import torchaudio


class ResidualBlock(nn.Module):
    """带通道扩展的残差块，保持时间维度"""

    def __init__(self, in_channels, out_channels, expansion=2):
        super().__init__()
        mid_channels = in_channels * expansion

        self.conv1 = nn.Conv1d(in_channels, mid_channels, kernel_size=3, padding='same', bias=False)
        self.ly1 = nn.LayerNorm(mid_channels)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv1d(mid_channels, out_channels, kernel_size=3, padding='same', bias=False)
        self.ln2 = nn.LayerNorm(out_channels)

        # 通道调整模块
        self.shortcut = (
            nn.Conv1d(in_channels, out_channels, 1)
            if in_channels != out_channels
            else nn.Identity()
        )

        self.final_activation = nn.ReLU()

    def forward(self, x):
        residual = self.shortcut(x.permute(0, 2, 1))
        out = x.permute(0, 2, 1)  # (B, T, in_dim) -> (B, in_dim, T)
        out = self.conv1(out)  # (B, in_dim, T) -> (B, mid_dim, T)
        out = out.permute(0, 2, 1)  # (B, mid_dim, T) -> (B, T, mid_dim)
        out = self.ly1(out)
        out = self.relu1(out)

        out = out.permute(0, 2, 1)  # (B, T, mid_dim) -> (B, mid_dim, T)
        out = self.conv2(out)  # (B, mid_dim, T) -> (B, out_dim, T)
        out = out.permute(0, 2, 1)  # (B, out_dim, T) -> (B, T, out_dim)
        out = self.ln2(out)

        out = out + residual.permute(0, 2, 1)

        return self.final_activation(out)


class ResNetFeatureExtractor(nn.Module):
    """ResNet特征提取器，支持直接输入原始音频并自动提取kaldi fbank特征"""

    def __init__(self, in_dim=80, out_dim=256, sample_rate=16000):
        super().__init__()
        self.sample_rate = sample_rate
        self.in_dim = in_dim

        # 初始通道扩展 (T, 80) -> (T, 128)
        dim_1 = 128  # 扩展的维度
        self.conv1 = nn.Conv1d(in_dim, dim_1, kernel_size=5, padding='same')
        self.ln1 = nn.LayerNorm(dim_1)
        self.relu1 = nn.ReLU()
        self.drop1 = nn.Dropout(0.1)

        # 残差块堆叠
        self.res_blocks = nn.Sequential(
            ResidualBlock(128, 128), ResidualBlock(128, 192), ResidualBlock(192, 256)
        )

        # 输出投影
        self.final_proj = nn.Conv1d(256, out_dim, kernel_size=1)

    def forward(self, x, x_lens):
        # 如果输入是原始音频 [B, 1, T]，先转为kaldi fbank
        if x.dim() == 3 and x.size(1) == 1:
            feats = []
            feat_lens = []
            for i in range(x.size(0)):
                wav = x[i, 0, :x_lens[i]].cpu()  # [T]
                fbank = torchaudio.compliance.kaldi.fbank(
                    wav.unsqueeze(0),
                    num_mel_bins=self.in_dim,
                    frame_length=25.0,
                    frame_shift=10.0,
                    sample_frequency=self.sample_rate,
                    dither=0.0,
                    energy_floor=0.0
                )
                # 保护：如果fbank帧数为0，补一帧全零
                if fbank.size(0) == 0:
                    fbank = torch.zeros(1, self.in_dim)
                feats.append(fbank)
                feat_lens.append(fbank.size(0))
            max_feat_len = max(feat_lens)
            feats_padded = torch.zeros(len(feats), max_feat_len, self.in_dim, device=x.device)
            for i, fbank in enumerate(feats):
                feats_padded[i, :fbank.size(0), :] = fbank.to(x.device)
            x = feats_padded # [B, 1, T] ->[B, T, 80]
            x_lens = torch.tensor(feat_lens, device=x.device)
        # 下面保持原有fbank特征处理流程
        x = x.permute(0, 2, 1)  # (B, T, 80) -> (B, 80, T)
        x = self.conv1(x)
        x = x.permute(0, 2, 1)
        x = self.ln1(x)
        x = self.relu1(x)
        x = self.drop1(x)
        x = self.res_blocks(x)
        x = x.permute(0, 2, 1)
        x = self.final_proj(x)
        return x.permute(0, 2, 1), x_lens  # (B, T, 256), x_lens


# 测试用例
if __name__ == "__main__":
    model = ResNetFeatureExtractor()
    dummy_input = torch.randn(32, 1, 16650)  # (B, 1, T)

    output, out_len = model(dummy_input, torch.tensor([16650]*32))

    print(
        f"输入形状: {dummy_input.shape} → 输出形状: {output.shape}"
    )  # 应输出 (32, 500, 256)
    print(out_len)
